tensor([[0.5700, 0.5600, 0.5700, 0.5600, 0.5700],
        [0.4300, 0.4400, 0.4300, 0.4400, 0.4300]], grad_fn=<RoundBackward1>)